package utils

import (
	"crypto/aes"
	"crypto/cipher"
	"fmt"
	"io"
	"math/rand"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// 加密配置
const (
	keySize        = 32        // AES-256
	nonceSize      = 16        // CTR模式可为16字节
	fileBufferSize = 64 * 1024 // 64KB分块
	maxWorkers     = 8         // 根据CPU核心数调整
)

// 加密文件结构: [Nonce][加密数据]
// EncryptFile 加密文件
//
// key: 加密密钥
// srcPath: 源文件路径
// dstPath: 目标文件路径（可选），如果为空则自动在源文件路径后追加 .enc
//
// 加密算法: AES-256 CTR 模式
//
// 返回错误信息
func EncryptFile(key []byte, srcPath, dstPath string) error {
	// 生成唯一Nonce
	nonce := make([]byte, nonceSize)
	if _, err := rand.Read(nonce); err != nil {
		return err
	}

	// 创建加密算法
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	stream := cipher.NewCTR(block, nonce) // CTR模式支持并行加密

	// 打开文件
	srcFile, err := os.OpenFile(srcPath, os.O_RDWR, 0666)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	// 写入Nonce到文件头部

	// 分块加密
	buf := make([]byte, fileBufferSize)
	write_buf := make([]byte, 0)

	for {
		n, err := srcFile.Read(buf)
		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 {
			break
		}

		// 原地加密（CTR模式特性）
		stream.XORKeyStream(buf[:n], buf[:n])
		write_buf = append(write_buf, buf[:n]...)
		fmt.Println("write_buf len:", len(write_buf))
	}

	dstPath = srcPath + ".enc"
	dstFile, err := os.OpenFile(dstPath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	defer dstFile.Close()

	if _, err = dstFile.Seek(0, 0); err != nil {
		return err
	}
	// 写入加密数据
	if _, err = dstFile.Write(nonce); err != nil {
		return err
	}

	if _, err = dstFile.Write(write_buf); err != nil {
		return err
	}
	return nil
}

// 解密（只需修改Nonce读取位置）
func DecryptFile(key []byte, srcPath, dstPath string) error {
	// 读取Nonce
	srcFile, err := os.OpenFile(srcPath, os.O_RDWR, 0666)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	// 读取Nonce
	nonce := make([]byte, nonceSize)
	if _, err = io.ReadFull(srcFile, nonce); err != nil {
		return err
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}

	stream := cipher.NewCTR(block, nonce) // CTR模式支持并行加密

	// 打开文件

	// 写入Nonce到文件头部

	// 分块加密
	buf := make([]byte, fileBufferSize)
	write_buf := make([]byte, 0)

	for {
		n, err := srcFile.Read(buf)
		if err != nil && err != io.EOF {
			return err
		}
		if n == 0 {
			break
		}

		// 原地加密（CTR模式特性）
		stream.XORKeyStream(buf[:n], buf[:n])
		write_buf = append(write_buf, buf[:n]...)
		fmt.Println("write_buf len:", len(write_buf))
	}
	if _, err = srcFile.Seek(0, 0); err != nil {
		return err
	}
	// 写入加密数据

	if _, err = srcFile.Write(write_buf); err != nil {
		return err
	}
	return nil

}

// 并行处理所有文件
func ProcessFiles(key []byte, rootDir string, outputDir string, operation_name string, operation func([]byte, string, string) error) error {
	start := time.Now()

	// 修改代码
	defer func() {
		end := time.Now()
		fmt.Printf("%s 耗时: %v\n", operation_name, end.Sub(start))
	}()

	filePaths := make(chan string, maxWorkers*2)
	errChan := make(chan error, maxWorkers)
	var wg sync.WaitGroup

	// 启动Worker池
	for i := 0; i < maxWorkers; i++ {

		wg.Add(1)
		go func() {
			defer wg.Done()
			for srcPath := range filePaths {
				relPath, _ := filepath.Rel(rootDir, srcPath)
				dstPath := filepath.Join(outputDir, relPath) // 示例输出目录

				// 创建目标目录
				if err := os.MkdirAll(filepath.Dir(dstPath), 0755); err != nil {
					errChan <- err
					continue
				}

				// 执行加密/解密
				if err := operation(key, srcPath, dstPath); err != nil {
					errChan <- fmt.Errorf("%s: %v", srcPath, err)
				}
			}
		}()
	}

	// 遍历目录并分发任务
	go func() {
		filepath.Walk(rootDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				errChan <- err
				return nil
			}
			if !info.IsDir() {
				filePaths <- path
			}
			return nil
		})
		defer close(filePaths)

	}()

	// 等待完成并收集错误
	go func() {
		wg.Wait()
		defer close(errChan)
	}()

	// 输出错误信息
	for err := range errChan {
		if err != nil {
			fmt.Printf("处理失败: %v\n", err)
		}
	}

	return nil
}

func GenerateRandomString(length int) string {

	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

	var seededRand *rand.Rand = rand.New(rand.NewSource(time.Now().UnixNano()))

	b := make([]byte, length)
	for i := range b {
		b[i] = charset[seededRand.Intn(len(charset))]
	}
	return string(b)
}
